#include "kiva/parser.h"

#include <algorithm>

namespace ast = pyast11;
namespace py = pybind11;

namespace kiva {
namespace {

class ParserImpl : public Parser {
 public:
  Expression ParseExpr(NodeRef node) override {
    NodeType node_type = node->type();
    DCHECK(Node::IsExpression(node_type));

    if (Node::IsBoolOp(node_type))
      return ParseBoolOp(node);
    if (Node::IsBinOp(node_type))
      return ParseBinOp(node);
    if (Node::IsUnaryOp(node_type))
      return ParseUnaryOp(node);
    if (Node::IsConstant(node_type))
      return ParseConstant(node);
    if (Node::IsCompare(node_type))
      return ParseCompare(node);

    switch (node_type) {
      case NodeType::kLambda:
        return ParseLambda(node);
      case NodeType::kTuple:
        return ParseTuple(node);
      case NodeType::kAwait:
        return ParseAwait(node);
      case NodeType::kYield:
        return ParseYield(node);
      case NodeType::kYieldFrom:
        return ParseYieldFrom(node);
      case NodeType::kName:
        return ParseName(node);
      case NodeType::kStarred:
        return ParseStarred(node);
      case NodeType::kAttribute:
        return ParseAttribute(node);
      case NodeType::kSubscript:
        return ParseSubscript(node);
      default:
        DCHECK(false) << "Unhandled NodeType" << static_cast<int>(node_type);
    }
  }

  Expression ParseSubscript(NodeRef node) {

  }

  Expression ParseAttribute(NodeRef node) {
    auto value = ParseExpr(node->child_at(0));
    auto attr = ast::MakeIdentifier(node->child_at(1)->string_value());
    return ast::Attribute(arena_, value, attr, context_);
  }

  Expression ParseStarred(NodeRef node) {
    auto value = ParseExpr(node->child_at(0));
    return ast::Starred(arena_, value, context_);
  }

  Expression ParseLambda(NodeRef node) {}
  Expression ParseYield(NodeRef node) {}
  Expression ParseYieldFrom(NodeRef node) {}

  Expression ParseAwait(NodeRef node) {
    auto value = ParseExpr(node->child_at(0));
    return ast::Await(arena_, value);
  }

  Expression ParseConstant(NodeRef node) {
    // return ast::Constant(arena_, py::str(""), node->object_value());
  }

  Expression ParseCompare(NodeRef node) {
    auto left = ParseExpr(node->child_at(0));
    int n_cmp = (node->num_children() - 1) / 2;
    ast::expr_list comparators(arena_, n_cmp);
    ast::list<ast::cmpop> ops(arena_, n_cmp);

    for (int i = 0; i < n_cmp; ++i) {
      auto op = node->child_at(2 * i + 2)->long_value();
      ops[i] = ToCompareOperator(static_cast<NodeType>(op));
      comparators[i] = ParseExpr(node->child_at(2 * i + 3));
    }
    return ast::Compare(arena_, left, ops, comparators);
  }

  Expression ParseBoolOp(NodeRef node) {
    DCHECK(Node::IsBoolOp(node->type()));
    BoolOperator op = ToBoolOperator(node->type());
    auto values = ast::MakeSeq(
        arena_, {ParseExpr(node->child_at(0)), ParseExpr(node->child_at(1))});
    return ast::BoolOp(arena_, op, values);
  }

  Expression ParseBinOp(NodeRef node) {
    Operator op = ToOperator(node->type());
    Expression lhs = ParseExpr(node->child_at(0));
    Expression rhs = ParseExpr(node->child_at(1));
    return ast::BinOp(arena_, lhs, op, rhs);
  }

  Expression ParseUnaryOp(NodeRef node) {
    UnaryOperator op = ToUnaryOperator(node->type());
    Expression operand = ParseExpr(node->child_at(0));
    return ast::UnaryOp(arena_, op, operand);
  }

  ast::keyword ParseKeyword(NodeRef node) {
    DCHECK(node->type() == NodeType::kAssign);
    DCHECK(node->num_children() == 2);
    DCHECK(node->child_at(0)->type() == NodeType::kName);
    auto id = ast::MakeIdentifier(node->child_at(0)->string_value());
    auto value = ParseExpr(node->child_at(1));
    return ast::keyword(arena_, std::move(id), value);
  }

  Expression ParseCall(NodeRef node) {
    auto func = ParseExpr(node->child_at(0));
    auto&& arguments = node->child_at(1)->children();
    int n_keywords = std::count_if(
        arguments.begin(), arguments.end(),
        [](NodeRef n) -> bool { return n->type() == NodeType::kAssign; });
    // should 0 size list be folded into null?
    ast::list<ast::keyword> keywords(arena_, n_keywords);
    ast::list<ast::expr> positionals(arena_, arguments.size() - n_keywords);

    int keyword_cnt = 0;
    int positional_cnt = 0;
    for (NodeRef child : arguments) {
      if (child->type() == NodeType::kAssign) {
        keywords[keyword_cnt++] = ParseKeyword(child);
      } else {
        positionals[positional_cnt++] = ParseExpr(child);
      }
    }
    return ast::Call(arena_, func, positionals, keywords);
  }

  ast::Name ParseName(NodeRef node) {
    return ast::Name(arena_, ast::MakeIdentifier(node->string_value()),
                     context_);
  }

  ast::arguments ParseArguments(NodeRef node) {
  }

  Statement ParseFunctionDef(NodeRef node) {
    auto node_type = node->type();
    DCHECK(node_type == NodeType::kFunctionDef ||
           node_type == NodeType::kAsyncFunctionDef);
    int num_children = node->num_children();
    DCHECK(num_children >= 2 && num_children <= 4);

    auto name = ast::MakeIdentifier(node->child_at(0)->string_value());
    auto arguments =
        ParseArguments(num_children > 2 ? node->child_at(1) : nullptr);
    auto returns =
        num_children == 4 ? ParseExpr(node->child_at(3)) : Expression{};
    auto body = ParseStatementList(node->child_at(num_children - 1));

    ast::list<ast::expr> decorators;
    if (!decorator_list_.empty())
      decorators =
          ast::MakeSeq(arena_, decorator_list_.begin(), decorator_list_.end());

    if (node_type == NodeType::kFunctionDef)
      return ast::FunctionDef(arena_, std::move(name), arguments, body,
                              decorators, returns);
    return ast::AsyncFunctionDef(arena_, std::move(name), arguments, body,
                                 decorators, returns);
  }

  Statement ParseReturn(NodeRef node) {
    auto value =
        node->num_children() == 1 ? ParseExpr(node->child_at(0)) : Expression();
    return ast::Return(arena_, value);
  }

  Statement ParseDelete(NodeRef node) {

  }

  Statement ParseClassDef(NodeRef node) {

  }

  Statement ParseAssign(NodeRef node) {
  }

  StatementList ParseStatementList(NodeRef node) override {
    auto iter = node->children().begin();
    auto end = node->children().end();
    std::vector<Statement> stmts;
    stmts.reserve(node->num_children());

    while (iter != end) {
      NodeRef cur(*iter);
      switch (cur->type()) {
        case NodeType::kFunctionDef:
        case NodeType::kAsyncFunctionDef:
          stmts.push_back(ParseFunctionDef(node));
          break;
        case NodeType::kClassDef:
          stmts.push_back(ParseClassDef(node));
          break;
        case NodeType::kReturn:
          stmts.push_back(ParseReturn(cur));
          break;
        case NodeType::kDelete:
          stmts.push_back(ParseDelete(cur));
          break;
        case NodeType::kAssign:
          stmts.push_back(ParseAssign(cur));
          break;


        default:
          break;
      }
    }
    return ast::MakeSeq(arena_, stmts.begin(), stmts.end());
  }

  Expression ParseTuple(NodeRef node) {
    DCHECK(node->type() == NodeType::kTuple);
    auto list = ast::list<ast::expr>(arena_, node->num_children());
    auto iter = list.begin();
    for (auto&& child : node->children())
      *iter = ParseExpr(child.get());
    return ast::Tuple(arena_, list, context_);
  }

 private:
  ast::expr_context context_ = ast::expr_context::Load;
  ast::Arena arena_;
  std::vector<ast::expr> decorator_list_;
};

template <typename T, NodeType First>
T EnumFromNodeType(NodeType type) {
  int val = static_cast<int>(type) - static_cast<int>(First);
  DCHECK(val >= 0);
  return static_cast<T>(val + 1);
}
}  // namespace

ast::operator_ Parser::ToOperator(NodeType type) {
  return EnumFromNodeType<Operator, NodeType::kBinaryExprFirst>(type);
}
ast::unaryop Parser::ToUnaryOperator(NodeType type) {
  return EnumFromNodeType<UnaryOperator, NodeType::kUnaryExprFirst>(type);
}
ast::boolop Parser::ToBoolOperator(NodeType type) {
  return EnumFromNodeType<BoolOperator, NodeType::kBoolOpFirst>(type);
}
ast::cmpop Parser::ToCompareOperator(NodeType type) {
  return EnumFromNodeType<CompareOperator, NodeType::kCompareFirst>(type);
}

std::unique_ptr<Parser> Parser::New() {
  return std::make_unique<ParserImpl>();
}

}  // namespace kiva
